# Copyright 2011 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Unit tests for ../appengine/excel_utils.py"""

import sys
sys.path.append('../appengine')

import StringIO
import unittest
import xlrd
import xlutils.filter
import xlwt

import excel_utils


class TestBook(object):
    """On-memory Excel workbook object for testing."""
    def __init__(self, book_contents):
        """Creates a test workbook.

        Example:
          book = TestBook(
            [{'name': 'Sheet1',
              'cells': [(0, 0, 'row0 col0'), (1, 2, 'row1 col2')]},
             {'name': 'Sheet2',
              'cells': [(3, 1, 'row3 col1')]}])

        Arguments:
          contents: List of dictionaries which represent worksheets. Each
                    worksheet dictionary has 'name' and 'cells' entries.
                    The value of 'name' is a string representing the sheet
                    name. The value of 'cells' is a list of tuples of row
                    index, column index and unicode string which represents
                    the cell's value.
        """
        book = xlwt.Workbook()
        for sheet in book_contents:
            new_sheet = book.add_sheet(sheet['name'])
            for (row, col, value) in sheet['cells']:
                new_sheet.write(row, col, value)
    
        buf = StringIO.StringIO()
        book.save(buf)
        self.__file_contents = buf.getvalue()
        self.__xlrd_book = xlrd.open_workbook(
            file_contents=self.__file_contents)
    
    def GetFileContents(self):
        """Gets the file contents of the workbook.

        Returns:
          A string representing the XLS file contents.
        """
        return self.__file_contents


    def GetXLRDBook(self):
        """Gets a XLRDBook object of the workbook.

        Returns:
          A xlrd.Book object.
        """
        return self.__xlrd_book


class ExcelUtilsTestBase(unittest.TestCase):
    def AssertEqualSheets(self, sheet1, sheet2, ignore_sheetname):
        """Tests if sheet1 is equivalent to sheet2.

        Arguments:
          sheet1: A xlrd.Sheet object.
          sheet2: A xlrd.Sheet object.
          ignore_sheetname: Boolean. If True, it doesn't test sheetname
                            matches.
        """
        if not ignore_sheetname:
            self.assertEqual(sheet1.name, sheet2.name)
        self.assertEqual(sheet1.nrows, sheet2.nrows)
        self.assertEqual(sheet1.ncols, sheet2.ncols)
        for row in range(sheet1.nrows):
            for col in range(sheet1.ncols):
                self.assertEqual(sheet1.cell_value(row, col),
                                 sheet2.cell_value(row, col))


    def AssertEqualBooks(self, book1, book2, ignore_sheetname=False):
        """Tests if book1 is equivalent to book2.

        Arguments:
          book1: A xlrd.Book object.
          book2: A xlrd.Book object.
          ignore_sheetname: Boolean. If True, it doesn't test sheetname
                            matches.
        """
        self.assertEqual(book1.nsheets, book2.nsheets)
        for i in range(book1.nsheets):
            sheet1 = book1.sheet_by_index(i)
            sheet2 = book2.sheet_by_index(i)
            self.AssertEqualSheets(sheet1, sheet2, ignore_sheetname)


class TestMultiWorkbookReader(ExcelUtilsTestBase):
    def testGetWorkbooks(self):
        book1 = TestBook([{'name': 'Sheet1',
                           'cells': [(0, 0, 'A'), (0, 1, 'B')]},
                          {'name': 'Sheet1-2',
                           'cells': [(1, 0, 'E'), (3, 1, 'F')]}])
        book2 = TestBook([{'name': 'Sheet2',
                           'cells': [(1, 0, 'C'), (2, 4, 'D')]}])
        books = [book1.GetXLRDBook(), book2.GetXLRDBook()]

        reader = excel_utils.MultiWorkbookReader(
            [book1.GetFileContents(), book2.GetFileContents()])
        readbooks = list(reader.get_workbooks())

        self.assertEqual(len(books), len(readbooks))
        for (book, (readbook, _)) in zip(books, readbooks):
            self.AssertEqualBooks(book, readbook)


class TestMergeWorkbookWriter(ExcelUtilsTestBase):
    def testMerge(self):
        book1 = TestBook([{'name': 'Sheet1-1',
                           'cells': [(0, 0, 'A'), (0, 1, 'B')]},
                          {'name': 'Sheet1-2',
                           'cells': [(1, 0, 'E'), (3, 1, 'F')]}])
        book2 = TestBook([{'name': 'Sheet2-1',
                           'cells': [(1, 0, 'C'), (2, 4, 'D')]}])
        books = [book1.GetXLRDBook(), book2.GetXLRDBook()]

        reader = excel_utils.MultiWorkbookReader(
            [book1.GetFileContents(), book2.GetFileContents()])
        writer = excel_utils.MergedWorkbookWriter()
        xlutils.filter.process(reader, writer)
        buf = StringIO.StringIO()
        writer.get_workbook().save(buf)
        mergedbook = xlrd.open_workbook(file_contents=buf.getvalue())

        self.assertEqual(sum(b.nsheets for b in books),
                         mergedbook.nsheets)
        book_index = 0
        sheet_index = 0
        for merged_index in range(mergedbook.nsheets):
            sheet = books[book_index].sheet_by_index(sheet_index)
            merged_sheet = mergedbook.sheet_by_index(merged_index)
            self.AssertEqualSheets(sheet, merged_sheet, ignore_sheetname=True)
            sheet_index += 1
            if sheet_index >= books[book_index].nsheets:
                book_index += 1
                sheet_index = 0


class TestFirstSheetFilter(ExcelUtilsTestBase):
    def testFilter(self):
        book = TestBook([{'name': 'Sheet1-1',
                          'cells': [(0, 0, 'A'), (0, 1, 'B')]},
                         {'name': 'Sheet1-2',
                          'cells': [(1, 0, 'E'), (3, 1, 'F')]}]).GetXLRDBook()

        reader = xlutils.filter.XLRDReader(book, 'test_first_sheet_filter')
        outbuf = StringIO.StringIO()
        writer = xlutils.filter.StreamWriter(outbuf)
        xlutils.filter.process(reader,
                               excel_utils.FirstSheetFilter(),
                               writer)
        filtered_book = xlrd.open_workbook(file_contents=outbuf.getvalue())

        self.assertEqual(1, filtered_book.nsheets)
        self.AssertEqualSheets(book.sheet_by_index(0),
                               filtered_book.sheet_by_index(0),
                               ignore_sheetname=False)

class TestMergeBooksFileter(ExcelUtilsTestBase):
    def testMerge(self):
        book1 = TestBook([{'name': 'Sheet1',
                           'cells': [(0, 0, 'A1'), (0, 1, 'B1')]},
                          {'name': 'Sheet22',
                           'cells': [(0, 0, 'C1'), (0, 1, 'D1')]}])
        book2 = TestBook([{'name': 'Sheet1',
                           'cells': [(0, 0, 'A2'), (0, 1, 'B2')]}])

        reader = excel_utils.MultiWorkbookReader(
            [book1.GetFileContents(), book2.GetFileContents()])
        outbuf = StringIO.StringIO()
        writer = xlutils.filter.StreamWriter(outbuf)
        xlutils.filter.process(reader,
                               excel_utils.MergeBooksFilter('merged'),
                               writer)
        merged_book = xlrd.open_workbook(file_contents=outbuf.getvalue())

        self.assertEqual(3, merged_book.nsheets)
        merged_sheet1 = merged_book.sheet_by_index(0)
        self.assertEqual('A1', merged_sheet1.cell_value(0, 0))
        self.assertEqual('B1', merged_sheet1.cell_value(0, 1))
        merged_sheet2 = merged_book.sheet_by_index(1)
        self.assertEqual('C1', merged_sheet2.cell_value(0, 0))
        self.assertEqual('D1', merged_sheet2.cell_value(0, 1))
        merged_sheet3 = merged_book.sheet_by_index(2)
        self.assertEqual('A2', merged_sheet3.cell_value(0, 0))
        self.assertEqual('B2', merged_sheet3.cell_value(0, 1))


class TestMergeSheetsFilter(ExcelUtilsTestBase):
    def testMerge(self):
        book = TestBook([{'name': 'Sheet1',
                          'cells': [(0, 0, 'A'), (0, 1, 'B')]},
                         {'name': 'Sheet2',
                          'cells': [(0, 0, 'C'), (0, 1, 'D')]}]).GetXLRDBook()
        
        reader = xlutils.filter.XLRDReader(book, 'test_merge_sheets_filter')
        outbuf = StringIO.StringIO()
        writer = xlutils.filter.StreamWriter(outbuf)
        xlutils.filter.process(reader,
                               excel_utils.MergeSheetsFilter('merged',
                                                             nrows_header=0),
                               writer)
        merged_book = xlrd.open_workbook(file_contents=outbuf.getvalue())

        self.assertEqual(1, merged_book.nsheets)
        merged_sheet = merged_book.sheet_by_index(0)
        self.assertEqual('merged', merged_sheet.name)
        self.assertEqual('A', merged_sheet.cell_value(0, 0))
        self.assertEqual('B', merged_sheet.cell_value(0, 1))
        self.assertEqual('C', merged_sheet.cell_value(1, 0))
        self.assertEqual('D', merged_sheet.cell_value(1, 1))

    def testMergeSkipHeaders(self):
        book = TestBook([{'name': 'Sheet1',
                          'cells': [(0, 0, 'header'),
                                    (1, 0, 'A'), (1, 1, 'B')]},
                         {'name': 'Sheet2',
                          'cells': [(0, 0, 'header'),
                                    (1, 0, 'C'), (1, 1, 'D')]}]).GetXLRDBook()
        
        reader = xlutils.filter.XLRDReader(book, 'test_merge_sheets_filter')
        outbuf = StringIO.StringIO()
        writer = xlutils.filter.StreamWriter(outbuf)
        xlutils.filter.process(reader,
                               excel_utils.MergeSheetsFilter('merged',
                                                             nrows_header=1),
                               writer)
        merged_book = xlrd.open_workbook(file_contents=outbuf.getvalue())

        self.assertEqual(1, merged_book.nsheets)
        merged_sheet = merged_book.sheet_by_index(0)
        self.assertEqual('merged', merged_sheet.name)
        self.assertEqual('header', merged_sheet.cell_value(0, 0))
        self.assertEqual('A', merged_sheet.cell_value(1, 0))
        self.assertEqual('B', merged_sheet.cell_value(1, 1))
        self.assertEqual('C', merged_sheet.cell_value(2, 0))
        self.assertEqual('D', merged_sheet.cell_value(2, 1))


class TestMergeFirstSheetsToOneSheet(ExcelUtilsTestBase):
    def testMerge(self):
        book1 = TestBook([{'name': 'Book1-1',
                           'cells': [(0, 0, 'A1'), (0, 1, 'B1')]},
                          {'name': 'Book1-2',
                           'cells': [(1, 0, 'C1'), (1, 1, 'D1')]}])
        book2 = TestBook([{'name': 'Book2-1',
                           'cells': [(0, 0, 'A2'), (0, 1, 'B2')]},
                          {'name': 'Book1-2',
                           'cells': [(1, 0, 'C2'), (1, 1, 'D2')]}])

        merged_file_contents = excel_utils.MergeFirstSheetsToOneSheet(
            [book1.GetFileContents(), book2.GetFileContents()],
            0, 'merged')
        merged_book = xlrd.open_workbook(file_contents=merged_file_contents)

        self.assertEqual(1, merged_book.nsheets)
        merged_sheet = merged_book.sheet_by_index(0)
        self.assertEqual('merged', merged_sheet.name)
        self.assertEqual('A1', merged_sheet.cell_value(0, 0))
        self.assertEqual('B1', merged_sheet.cell_value(0, 1))
        self.assertEqual('A2', merged_sheet.cell_value(1, 0))
        self.assertEqual('B2', merged_sheet.cell_value(1, 1))


class TestMergeAllSheetsToOneSheet(ExcelUtilsTestBase):
    def testMerge(self):
        book1 = TestBook([{'name': 'Sheet1',
                           'cells': [(0, 0, 'A1'), (0, 1, 'B1')]},
                          {'name': 'Sheet2',
                           'cells': [(0, 0, 'C1'), (0, 1, 'D1')]}])
        book2 = TestBook([{'name': 'Sheet1',
                           'cells': [(0, 0, 'A2'), (0, 1, 'B2')]},
                          {'name': 'Sheet2',
                           'cells': [(0, 0, 'C2'), (0, 1, 'D2')]}])

        merged_file_contents = excel_utils.MergeAllSheetsToOneSheet(
            [book1.GetFileContents(), book2.GetFileContents()],
            0, 'merged')
        merged_book = xlrd.open_workbook(file_contents=merged_file_contents)

        self.assertEqual(1, merged_book.nsheets)
        merged_sheet = merged_book.sheet_by_index(0)
        self.assertEqual('merged', merged_sheet.name)
        self.assertEqual('A1', merged_sheet.cell_value(0, 0))
        self.assertEqual('B1', merged_sheet.cell_value(0, 1))
        self.assertEqual('C1', merged_sheet.cell_value(1, 0))
        self.assertEqual('D1', merged_sheet.cell_value(1, 1))
        self.assertEqual('A2', merged_sheet.cell_value(2, 0))
        self.assertEqual('B2', merged_sheet.cell_value(2, 1))
        self.assertEqual('C2', merged_sheet.cell_value(3, 0))
        self.assertEqual('D2', merged_sheet.cell_value(3, 1))


if __name__ == '__main__':
    unittest.main()
